import argparse
import glob
import json
import os
import re

import pandas as pd

from iflytech_assistant.assistant.dataclasses import FiveSuggestions


def preprocess(target: str, tag: str, user_input: str):
    tag = re.sub(r"[^\u4e00-\u9fa5a-zA-Z]", "", tag)
    target = target.replace("发", "")
    return {
        "target": target.strip(),
        "tag": tag.strip(),
        "user_input": user_input.strip(),
    }


user_prompt = """
我现在想学会跟别人高情商沟通。并且有一个要求，需要伪装成{preferenced_style}的人设。

现在我在微信聊天框里面，对面是我的{target}，我想要表达的话是：“{user_input}”

帮我分析分析我目前的情况，应该怎么表达比较好呢？我想让你先帮我分析下。如果我的意图不是很明确，也可以帮我补充补充。
口语一点，不要太书面，我要直接发到对方微信上的。你给我整的太死板书面了，我没法直接发过去。
分析完事后，麻烦你给我五条不同的改写结果，1. 2. 3. 4. 5.，每条改写结果之间用空行隔开。
"""

# 单轮
template = {"param": {"max_tokens": 2048}, "id": "h_001", "query": ""}


# 多轮
{
    "q_a_history": [
        {"role": "user", "content": ""},
        {"role": "assistant", "content": ""},
    ],
    "query": ["q1", "q2"],
    "id": "",
    "param": {"temperature": 0},
}
parser = argparse.ArgumentParser()
parser.add_argument(
    "-n",
    "--name",
    type=str,
)
args = parser.parse_args()

df = pd.read_excel(".vscode/高情商沟通cot标注/extended_data2.xlsx")

tasks = []
for i, row in df.iterrows():
    user_input = row["用户输入"]
    target = row["沟通对象"]
    preferenced_style = row["用户偏好"]

    result = preprocess(target, preferenced_style, user_input)
    target = result["target"]
    preferenced_style = result["tag"]
    user_input = result["user_input"]

    up = user_prompt.format(
        user_input=user_input.strip(),
        target=target,
        preferenced_style=preferenced_style,
    )
    query = up

    line = template.copy()
    line["query"] = query
    tasks.append(line)

df = pd.DataFrame(tasks)
df.to_excel(".vscode/test2.xlsx", index=False)
# to jsonl
df.to_json(".vscode/test2.jsonl", orient="records", lines=True, force_ascii=False)


# for file in glob.glob(".vscode/高情商沟通标注/*.jsonl"):
#     filename = os.path.basename(file)
#     lines = []
#
#     with open(file, "r") as f:
#         for line in f:
#             line = json.loads(line)
#             answer = line["answer"]
#             suggestions = FiveSuggestions.from_llm_response(answer)
#             suggestions.suggestions = suggestions.suggestions[-5:]
#             lines.append(str(suggestions))
#         df[filename] = lines
#
# df.to_excel(".vscode/高情商沟通标注/标注结果.xlsx", index=False)
